# -- coding: utf-8 --
import os
from tqdm import tqdm

from sklearn.metrics import classification_report

os.environ["CUDA_VISIBLE_DEVICES"] = '1'
import torch
import json
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, TaskType, get_peft_model

torch.set_default_tensor_type(torch.cuda.HalfTensor)

model_path = '/workspace/psycho_trainning/resources/Qwen-14B-Chat'
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device_map='auto')

peft_path = '/workspace/psycho_trainning/ckpts/psycho/lora/only_modified_with_states_0531/checkpoint-600/adapter_model.bin'

peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    inference_mode=True,
    r=8,
    target_modules=['c_attn'],
    lora_alpha=32,
    lora_dropout=0.1,
)

model = get_peft_model(model, peft_config)
model.load_state_dict(torch.load(peft_path), strict=False)
model.eval()

text = "##角色与目标##\n你是优秀的心理咨询师，你擅长焦点解决短期治疗技术、富有同理心和共情能力。你的目标是发挥你的专业性，深度共情来访者、安慰来访者，为来访者提供情绪价值、缓解来访者的负面情绪。同时在必要的时候，通过焦点解决短期治疗技术，找到用户的需求和目标，帮助用户解决问题。另外，当用户需要明确的建议时，提供专业的心理教育。\n##说明##\n请结合一段json格式的输入数据，数据内容为来访者和咨询师的对话内容，和其他的一些参考信息，具体内容如下： \n\"history\"项包含一段来访者和心理咨询师的历史对话， \"history_actions\"是已经做过的历史动作，\"role\"项是心理咨询师，\"prev_global_states\"项包含历史对话中来访者的历史状态。\n 输入的数据如下：{'history': '来访者：我做了个噩梦\\n', 'history_actions': [], 'role': '咨询师', 'prev_global_states': ''}\n ##要求##\n现在你作为心理咨询师，对来访者当前的对话进行回应。\n请先以json格式输出\"来访者四象限\"(请根据\"history\"中来访者最后一句回复分析出)，\"咨询师四象限\"，\"流程\"，\"来访者想要的\"，\"对来访者来说重要的\"，\"来访者具备的能力/资源/优势\"，\n再根据以上的信息，输出\"动作\"和\"内容\n##澄清##\n\"动作回复\"中的动作尽量避免包含历史动作，输出的回复也不要跟历史回复相似或重复\n输出如下：\n"


ids = tokenizer.encode(text)
input_ids = torch.LongTensor([ids]).cuda()
output = model.generate(
    input_ids=input_ids,
    temperature=0.5
)

out_text = tokenizer.decode(output[0].cpu().numpy().tolist(), skip_special_tokens=True)
print()

new_base_model = model.merge_and_unload()
new_base_model.save_pretrained('/workspace/psycho_trainning/ckpts/psycho/datacoach/qwen_0531')
tokenizer.save_pretrained('/workspace/psycho_trainning/ckpts/psycho/datacoach/qwen_0531')